import os
import numpy
import pandas as pd
import time
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


file_tmp1 = open("./ranked_0.pdb", 'r')
file_tmp2 = open("./ranked_1.pdb", 'r')
file_tmp3 = open("./ranked_2.pdb", 'r')
file_tmp4 = open("./ranked_3.pdb", 'r')
file_tmp5 = open("./ranked_4.pdb", 'r')
tmp_lines1 = [i for i in file_tmp1]
tmp_lines2 = [i for i in file_tmp2]
tmp_lines3 = [i for i in file_tmp3]
tmp_lines4 = [i for i in file_tmp4]
tmp_lines5 = [i for i in file_tmp5]
file_tmp1.close()
file_tmp2.close()
file_tmp3.close()
file_tmp4.close()
file_tmp5.close()
tmp_lines1_clean = []
tmp_lines2_clean = []
tmp_lines3_clean = []
tmp_lines4_clean = []
tmp_lines5_clean = []

for i in tmp_lines1:
    if i[0:4] == "ATOM":
        tmp_lines1_clean.append(i)

for i in tmp_lines2:
    if i[0:4] == "ATOM":
        tmp_lines2_clean.append(i)

for i in tmp_lines3:
    if i[0:4] == "ATOM":
        tmp_lines3_clean.append(i)

for i in tmp_lines4:
    if i[0:4] == "ATOM":
        tmp_lines4_clean.append(i)

for i in tmp_lines5:
    if i[0:4] == "ATOM":
        tmp_lines5_clean.append(i)


file_tmp1 = open("./ranked_0_clean.pdb", 'w+')
file_tmp1.writelines(tmp_lines1_clean)
file_tmp1.close()
file_tmp2 = open("./ranked_1_clean.pdb", 'w+')
file_tmp2.writelines(tmp_lines2_clean)
file_tmp2.close()
file_tmp3 = open("./ranked_2_clean.pdb", 'w+')
file_tmp3.writelines(tmp_lines3_clean)
file_tmp3.close()
file_tmp4 = open("./ranked_3_clean.pdb", 'w+')
file_tmp4.writelines(tmp_lines4_clean)
file_tmp4.close()
file_tmp5 = open("./ranked_4_clean.pdb", 'w+')
file_tmp5.writelines(tmp_lines5_clean)
file_tmp5.close()

file_tmp1 = open("./ranked_0_clean.pdb", 'r')
file_tmp2 = open("./ranked_1_clean.pdb", 'r')
file_tmp3 = open("./ranked_2_clean.pdb", 'r')
file_tmp4 = open("./ranked_3_clean.pdb", 'r')
file_tmp5 = open("./ranked_4_clean.pdb", 'r')
tmp_lines1 = [i for i in file_tmp1]
tmp_lines2 = [i for i in file_tmp2]
tmp_lines3 = [i for i in file_tmp3]
tmp_lines4 = [i for i in file_tmp4]
tmp_lines5 = [i for i in file_tmp5]

file_tmp1.close()
file_tmp2.close()
file_tmp3.close()
file_tmp4.close()
file_tmp5.close()

aa_number = []
plddt1 = []
plddt2 = []
plddt3 = []
plddt4 = []
plddt5 = []


for i in range(len(tmp_lines1)):
    if int(tmp_lines1[i][22:30]) not in aa_number:
        aa_number.append(int(tmp_lines1[i][22:30]))
        plddt1.append(float(tmp_lines1[i][61:66]))
aa_number = []
for i in range(len(tmp_lines2)):
    if int(tmp_lines2[i][22:30]) not in aa_number:
        aa_number.append(int(tmp_lines2[i][22:30]))
        plddt2.append(float(tmp_lines2[i][61:66]))
aa_number = []
for i in range(len(tmp_lines3)):
    if int(tmp_lines3[i][22:30]) not in aa_number:
        aa_number.append(int(tmp_lines3[i][22:30]))
        plddt3.append(float(tmp_lines3[i][61:66]))
aa_number = []
for i in range(len(tmp_lines4)):
    if int(tmp_lines4[i][22:30]) not in aa_number:
        aa_number.append(int(tmp_lines4[i][22:30]))
        plddt4.append(float(tmp_lines4[i][61:66]))
aa_number = []
for i in range(len(tmp_lines5)):
    if int(tmp_lines5[i][22:30]) not in aa_number:
        aa_number.append(int(tmp_lines5[i][22:30]))
        plddt5.append(float(tmp_lines5[i][61:66]))
dict1 = {'aa_number':aa_number,'plddt':plddt1}
dict2 = {'aa_number':aa_number,'plddt':plddt2}
dict3 = {'aa_number':aa_number,'plddt':plddt3}
dict4 = {'aa_number':aa_number,'plddt':plddt4}
dict5 = {'aa_number':aa_number,'plddt':plddt5}


df1 = pd.DataFrame(dict1)
df2 = pd.DataFrame(dict2)
df3 = pd.DataFrame(dict3)
df4 = pd.DataFrame(dict4)
df5 = pd.DataFrame(dict5)
print(df1)
print(df2)
print(df3)
print(df4)
print(df5)


x = df1['aa_number']
y1 = df1['plddt']
y2 = df2['plddt']
y3 = df3['plddt']
y4 = df4['plddt']
y5 = df5['plddt']

# Plot Line1 (Left Y Axis)
fig, ax1 = plt.subplots(1,1,figsize=(16,9), dpi= 600)

ax1.plot(x, y5, color='tab:orange',label="ranked_4")
ax1.plot(x, y4, color='tab:purple',label="ranked_3")
ax1.plot(x, y3, color='tab:blue'  ,label="ranked_2")
ax1.plot(x, y2, color='tab:green' ,label="ranked_1")
ax1.plot(x, y1, color='tab:red'   ,label="ranked_0")
plt.legend()
# Decorations
# ax1 (left Y axis)
ax1.set_xlabel('Amino acid number', fontsize=24)
# ax1.tick_params(axis='x', rotation=0, labelsize=10,labelcolor='tab:red')
ax1.tick_params(axis='x', rotation=0, labelsize=10)
#显示所有横坐标
# plt.xticks(x)
ax1.xaxis.set_major_locator(ticker.MultipleLocator(base=30))

ax1.set_ylabel('Plddt(0~100)', fontsize=24)
# ax1.tick_params(axis='y', rotation=0, labelcolor='tab:red' )
ax1.tick_params(axis='y', rotation=0)
ax1.grid(alpha=.1)


fig.tight_layout()

plt.savefig('./ranked_merge.png')
# plt.show()

for i in range(5):
    if os.path.exists("./ranked_{}_clean.pdb".format(i)):
        os.remove("./ranked_{}_clean.pdb".format(i))
    else:
        print("The file ranked_{}_clean.pdb does not exist".format(i))






